import os from 'os';
import { AsyncQueue } from "../common/AsyncQueue";

 
const cpuCount = os.cpus().length;

type Gene = number[];

interface GeneticAlgorithmOptions {
  /** 种群大小: 代表每一代中遗传算法的个体数量 */
  populationSize: number;
  /** 突变率 */
  mutationRate: number;
  /** 交叉率 */
  crossoverRate: number;
  /** 迭代代数 */
  generations: number;
  parameterRanges: [number, number][];
  parameterStep: number;
  fitnessFunction: (genes: Gene) => Promise<number>;
}

export class GeneticAlgorithm {
  private population: Gene[];
  private options: GeneticAlgorithmOptions;

  constructor(options: GeneticAlgorithmOptions) {
    this.options = options;
    this.population = [];

    // Initialize the population with random genes
    for (let i = 0; i < this.options.populationSize; i++) {
      const genes = this.getRandomGenes();
      this.population.push(genes);
    }
  }

  public async run(): Promise<Gene> {
    let fittest = this.population[0];

    for (
      let generation = 0;
      generation < this.options.generations;
      generation++
    ) {
      // Evaluate the fitness of each gene
      // const fitnessScores = this.population.map(genes => ({
      //   genes,
      //   fitness: this.options.fitnessFunction(genes)
      // }));
      const fitnessScores: {genes: Gene, fitness: number}[] = [];
      const fitnessAsyncQueue = new AsyncQueue(cpuCount - 1);

      for (let i = 0; i < this.population.length; i++) {
        const genes = this.population[i];
        fitnessScores.push({
          genes,
          fitness: 0,
        })
        // 提高并发，使用异步队列
        fitnessAsyncQueue.enqueue(() => {
          return this.options.fitnessFunction(genes)
        });
      }
      const fitness = await fitnessAsyncQueue.awaitAll();
      fitness.forEach((score, index) => {
        fitnessScores[index].fitness = score
      });
      // Sort the population by fitness
      fitnessScores.sort((a, b) => b.fitness - a.fitness);

      // Select the fittest gene
      const currentGenes = fitnessScores[0].genes;
      const [ currentFittest, nextFittest ] = await Promise.all([
        this.options.fitnessFunction(currentGenes),
        this.options.fitnessFunction(fittest),
      ]);
      
      if (
        currentFittest >
        nextFittest
      ) {
        fittest = currentGenes;
      }

      // Create the next generation
      const nextGeneration: Gene[] = [];
      while (nextGeneration.length < this.options.populationSize) {
        // Select two parent genes
        const parent1 = this.selectParent(fitnessScores);
        const parent2 = this.selectParent(fitnessScores);

        // Crossover the parents' genes
        const child = this.crossover(parent1, parent2);

        // Mutate the child's genes
        this.mutate(child);

        // Add the child to the next generation
        nextGeneration.push(child);
      }

      this.population = nextGeneration;
    }

    return fittest;
  }

  private getRandomGenes(): Gene {
    const genes: Gene = [];
    for (const [min, max] of this.options.parameterRanges) {
      const range = Math.round((max - min) / this.options.parameterStep);
      const value =
        min + Math.random() * range * this.options.parameterStep;
      genes.push(value);
    }
    return genes;
  }

  private selectParent(
    fitnessScores: { genes: Gene; fitness: number }[]
  ): Gene {
    const totalFitness = fitnessScores.reduce(
      (sum, { fitness }) => sum + fitness,
      0
    );
    let cumulativeFitness = 0;
    const threshold = Math.random() * totalFitness;
    for (const { genes, fitness } of fitnessScores) {
      cumulativeFitness += fitness;
      if (cumulativeFitness >= threshold) {
        return genes;
      }
    }
    return fitnessScores[fitnessScores.length - 1].genes;
  }

  private crossover(parent1: Gene, parent2: Gene): Gene {
    const child: Gene = [];
    for (let i = 0; i < parent1.length; i++) {
      if (Math.random() < this.options.crossoverRate) {
        child.push(parent1[i]);
      } else {
        child.push(parent2[i]);
      }
    }
    return child;
  }

  private mutate(genes: Gene): void {
    for (let i = 0; i < genes.length; i++) {
      if (Math.random() < this.options.mutationRate) {
        const [min, max] = this.options.parameterRanges[i];
        const range = Math.round(
          (max - min) / this.options.parameterStep
        );
        const mutation =
          (Math.random() - 0.5) * range * this.options.parameterStep;
        genes[i] += mutation;
        genes[i] = Math.max(min, Math.min(max, genes[i]));
      }
    }
  }
}

